from A000_计算中国的资本存量 import CapitalStockData;
from utils import prepare_Y
import pandas as pd


def cal_ratio_bdy_and_edy(cap, vad, bdy, edy):
    cap.periodically_add_k(build_discount_year=bdy, equipment_discount_year=edy)
    # 根据建筑安装工程价格指数和设备购置价格指数，计算2024年价格调整后的资本存量
    cap.reset_price_of_k()
    # cap.save_dic() # 保存2024年价格调整后的资本存量
    df = cap.get_k_df()
    vad = vad.loc[df.index]

    new_cols = list(range(2004,2024))
    # 行列对齐
    vad = vad[new_cols]
    df = df[new_cols]
    ratio = vad.div(df)
    return ratio



cap = CapitalStockData()
# 先递推2018年之前的分行业固定资本形成
dddf, correct_data = cap.tab_build_and_equipment_before2018()
# 再递推2018-2024年分行业固定资本形成
build_2024, equipment_2024 = cap.tab_build_and_equipment_after2018()
# 然后设定建筑折旧年限55年，设备折旧年限16年，计算资本存量


if __name__ == '__main__':
    vad = prepare_Y() # 获取每个行业的增加值数据

    ratio_values = []
    df_by_industry_dic = {}

    sample_df = cal_ratio_bdy_and_edy(cap, vad, 55, 10)

    for bdy in range(30,56):
        for edy in range(5,15):
            print(bdy, edy)
            ratio = cal_ratio_bdy_and_edy(cap, vad, bdy, edy)
            for idx in ratio.index:
                for col in ratio.columns:
                    ratio_values.append(
                        {
                            'bdy':bdy, 
                            'edy':edy,
                            'industry':idx,
                            'year':col,
                            'ratio':ratio.loc[idx, col]
                        })

    all_df = pd.DataFrame(ratio_values)

